import gc
import glob
import random
import torch
import re

TRAIN=1
if TRAIN:
    test_id=str(5)
    #mode='test'
    mode = 'train'
    #data_path='../data/cnndm_data1/cnndm.test.'+test_id+'.pt'#测试集数据集
    data_path = '../data/cnndm_data1/cnndm.'+mode+'.' + test_id + '.pt'
    dataset=torch.load(data_path)
    #只保留前10个dict的句子
    cnt=0
    test_sent=[]
    for dict in dataset:
        if cnt>=10:
            break
        text_list=dict['src_txt']
        #这里我只保留前100句话
        for line in  text_list:
            test_sent.append(line+'\n')
        cnt+=1
    #先构建新的dict存储
    dataset=dataset[:10]
    torch.save(dataset,'cnndm_'+mode+'.'+test_id+'.pt')

    #以txt形式保存 方便之后supert评测
    f1=open('cnndm_'+mode+'_'+test_id+'.txt','w',encoding='utf-8')
    for sent in test_sent:
        f1.write(sent)
    f1.close()

#下面是处理summary  将<q>的分隔符转换为,
#预处理部分不需要执行这个， 完成摘要生成才需要
if TRAIN!=1:
    sum_id=str(5)
    summary_path='../results/result_'+sum_id+'_gold.txt'
    fread=open(summary_path,'r')
    summary=fread.readlines()
    new_sum=[]
    pattern='<q>'
    len_pattern=len(pattern)
    for i in range(len(summary)):
        new_line=summary[i].rstrip('\n')
        while new_line.find(pattern)!=-1:
            ind=new_line.find(pattern)
            new_line=new_line[:ind]+','+new_line[ind+len_pattern:]
        new_sum.append(new_line)

    summary=' '.join(new_sum)
    f2=open('summary_test_'+sum_id+'.txt','w')
    f2.write(summary)
    f2.close()